#include "tools.cpp"

class Solution_simple {
  public:
    ListNode* mergeKLists(vector<ListNode*>& lists) {
        auto head = new ListNode(1);
        mergeKLists(lists, head);
        auto tmp = head->next;
        delete head;
        return tmp;
    }

    void mergeKLists(vector<ListNode*>& lists, ListNode* head) {
        int count = 0;
        int min_index = 0;
        for (int i = 0; i < lists.size(); ++i) {
            if (lists[i] != nullptr) {
                ++count;
                if (lists[min_index] == nullptr) min_index = i;
                if (lists[i]->val < lists[min_index]->val) min_index = i;
            }
        }
        if (count == 0) return;

        head->next = lists[min_index];
        lists[min_index] = lists[min_index]->next;
        mergeKLists(lists, head->next);
    }
};

class Solution {
  public:
    ListNode* mergeKLists(vector<ListNode*>& lists) {
        std::function<bool(ListNode*, ListNode*)> compare = [](ListNode* rhs, ListNode* lhs) {
            // if (not lhs) return true;
            // if (not rhs) return false;
            return rhs->val > lhs->val;
        };
        std::priority_queue<ListNode*, std::vector<ListNode*>, decltype(compare)> small_heap(compare);
        for (auto list : lists)
            if (list != nullptr) small_heap.push(list);

        auto head = new ListNode(0);
        mergeKLists(small_heap, head);

        auto tmp = head->next;
        head->next = nullptr;
        delete head;
        return tmp;
    }

    void mergeKLists(
        std::priority_queue<ListNode*, std::vector<ListNode*>, std::function<bool(ListNode*, ListNode*)>>& small_heap,
        ListNode* pre) {
        // if (small_heap.empty()) return;
        // pre->next = small_heap.top();
        // small_heap.pop();
        // if (pre->next->next != nullptr) small_heap.push(pre->next->next);
        // mergeKLists(small_heap, pre->next);
        while (not small_heap.empty()) {
            pre->next = small_heap.top();
            small_heap.pop();
            if (pre->next->next != nullptr) small_heap.push(pre->next->next);
            pre = pre->next;
        }
    }
};

int main() {
    Solution s;
    auto head1 = creat_list(std::vector<int>{1, 4, 5});
    auto head2 = creat_list(std::vector<int>{1, 3, 4});
    auto head3 = creat_list(std::vector<int>{2, 6});
    std::vector<ListNode*> lists{head1, head2, head3};
    auto head = s.mergeKLists(lists);
    print(head);
}
